/*
 * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
 * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
 */

package org.jetbrains.kotlin.cpp

import org.gradle.api.DefaultTask
import org.gradle.api.file.*
import org.gradle.api.internal.file.FileOperations
import org.gradle.api.model.ObjectFactory
import org.gradle.api.provider.ListProperty
import org.gradle.api.provider.MapProperty
import org.gradle.api.provider.Property
import org.gradle.api.tasks.*
import org.gradle.process.ExecOperations
import org.gradle.workers.WorkAction
import org.gradle.workers.WorkParameters
import org.gradle.workers.WorkerExecutor
import org.jetbrains.kotlin.bitcode.CompileToBitcodePlugin
import org.jetbrains.kotlin.clangArgs
import org.jetbrains.kotlin.execLlvmUtility
import org.jetbrains.kotlin.konan.target.PlatformManager
import org.jetbrains.kotlin.platformManagerProvider
import java.io.File
import javax.inject.Inject

private abstract class ClangFrontendJob : WorkAction<ClangFrontendJob.Parameters> {
    interface Parameters : WorkParameters {
        val workingDirectory: DirectoryProperty
        val targetName: Property<String>
        val inputPathRelativeToWorkingDir: Property<String>
        val outputFile: RegularFileProperty
        val compilerExecutable: Property<String>
        val arguments: ListProperty<String>
        val platformManager: Property<PlatformManager>
        val clangPaths: Property<String>
    }

    @get:Inject
    abstract val objects: ObjectFactory

    @get:Inject
    abstract val execOperations: ExecOperations

    override fun execute() {
        with(parameters) {
            outputFile.get().asFile.parentFile.mkdirs()
            execOperations.execLlvmUtility(platformManager.get(), compilerExecutable.get()) {
                workingDir = workingDirectory.asFile.get()
                args = arguments.get() + listOf(inputPathRelativeToWorkingDir.get(), "-o", outputFile.get().asFile.absolutePath)
                environment["PATH"] = clangPaths.get() + File.pathSeparator + environment["PATH"]
            }
        }
    }
}

/**
 * Compiling [inputFiles] with clang into LLVM bitcode in [outputFiles].
 *
 * @see CompileToBitcodePlugin
 */
@CacheableTask
open class ClangFrontend @Inject constructor(
        objects: ObjectFactory,
        private val workerExecutor: WorkerExecutor,
        private val layout: ProjectLayout,
        private val fileOperations: FileOperations,
) : DefaultTask() {
    protected data class WorkUnit(
            /**
             * Input file for [compiler].
             */
            @get:Input
            val inputPathRelativeToWorkingDir: String,

            /**
             * Output file produced by [compiler] for [inputPathRelativeToWorkingDir]
             */
            @get:OutputFile
            val outputFile: RegularFile,
    )

    /**
     * Where to put bitcode files generated by clang.
     */
    @get:Internal("Used to compute workUnits")
    val outputDirectory: DirectoryProperty = objects.directoryProperty()

    /**
     * Source files to compile from.
     */
    @get:SkipWhenEmpty
    @get:InputFiles
    @get:PathSensitive(PathSensitivity.NONE) // manually computed: [workUnits]
    val inputFiles: ConfigurableFileCollection = objects.fileCollection()

    /**
     * Will select the appropriate compiler and additional flags.
     */
    // TODO: Consider specifying full clang execution here and configure it from the plugin.
    @get:Input
    val targetName: Property<String> = objects.property(String::class.java)

    /**
     * The compiler to be used.
     *
     * Currently only `clang` and `clang++` are supported.
     */
    @get:Input
    val compiler: Property<String> = objects.property(String::class.java)

    /**
     * Extra arguments for [compiler].
     */
    @get:Input
    val arguments: ListProperty<String> = objects.listProperty(String::class.java)

    /**
     * Working directory for [compiler].
     *
     * All inputs will be passed to the compiler as relative paths to this directory.
     */
    @get:Internal("Used to compute workUnits and headersPathsRelativeToWorkingDir")
    val workingDirectory: DirectoryProperty = objects.directoryProperty()

    @get:Internal
    val reproducibilityRootsMap: MapProperty<File, String> = objects.mapProperty(File::class.java, String::class.java)

    /**
     * Locations to search for headers.
     *
     * Will be passed to the compiler as `-I…` and will also be used to compute task dependencies: recompile if the headers change.
     */
    @get:Nested
    val headersDirs: CppHeadersSet = objects.cppHeadersSet().apply {
        workingDir.set(this@ClangFrontend.workingDirectory)
    }

    /**
     * [WorkUnit]s for [compiler].
     *
     * For each [WorkUnit] one [compiler] invocation will be executed.
     */
    @get:Nested
    protected val workUnits = workingDirectory.zip(outputDirectory) { l, r -> l to r }.zip(inputFiles.elements) { (base, out), files ->
        files.map {
            val relativePath = it.asFile.toRelativeString(base.asFile)
            WorkUnit(relativePath, out.file(relativePath.replaceAfterLast(".", "bc")))
        }
    }

    @get:Nested
    protected val platformManagerProvider = objects.platformManagerProvider(project)

    @TaskAction
    fun compile() {
        val workQueue = workerExecutor.noIsolation()

        val platformManager = platformManagerProvider.platformManager.get()
        val target = platformManager.targetByName(targetName.get())
        val compilerSpecificArgs = platformManager.clangArgs(target, compiler.get())
        val clangPaths = fileOperations.configurableFiles(platformManager.hostPlatform.clang.clangPaths).asPath

        workUnits.get().forEach { workUnit ->
            workQueue.submit(ClangFrontendJob::class.java) {
                workingDirectory.set(this@ClangFrontend.workingDirectory)
                targetName.set(this@ClangFrontend.targetName)
                inputPathRelativeToWorkingDir.set(workUnit.inputPathRelativeToWorkingDir)
                outputFile.set(workUnit.outputFile)
                compilerExecutable.set(this@ClangFrontend.compiler)
                arguments.set(defaultCompilerFlags(this@ClangFrontend.headersDirs, this@ClangFrontend.reproducibilityRootsMap.get()))
                arguments.addAll(compilerSpecificArgs)
                arguments.addAll(this@ClangFrontend.arguments)
                this.platformManager.set(platformManager)
                this.clangPaths.set(clangPaths)
            }
        }
    }

    companion object {
        internal fun defaultCompilerFlags(headersDirs: CppHeadersSet, reproducibilityRootsMap: Map<File, String>): List<String> = buildList {
            add("-c")
            add("-emit-llvm")
            addAll(headersDirs.asCompilerArguments.get())
            // Prevent generated binaries from containing absolute paths
            addAll(reproducibilityRootsMap.map {
                "-ffile-prefix-map=${it.key}=${it.value}"
            })
        }
    }
}
